package avlTree;

public class AVLTree {
    public static class TreeNode{
        public int val;
        public int bf = 0;
        public TreeNode left;
        public TreeNode right;
        public TreeNode parent;

        public TreeNode(int val) {
            this.val = val;
        }

        @Override
        public String toString() {
            return "TreeNode{" +
                    "val=" + val +
                    ", bf=" + bf +
                    '}';
        }
    }

    public TreeNode root = null;

    public void insert(int val) {
        TreeNode node = new TreeNode(val);
        if (root == null) {
            root = node;
            return;
        }
        TreeNode parent = null;
        TreeNode tmp = root;
        while (tmp != null) {
            if (tmp.val < node.val) {
                parent = tmp;
                tmp = tmp.right;
            } else if (tmp.val > node.val) {
                parent = tmp;
                tmp = tmp.left;
            } else {
                return;
            }
        }
        if (parent.val < val) {
            parent.right = node;
            parent.bf += 1;
        } else {
            parent.left = node;
            parent.bf -= 1;
        }
        node.parent = parent;
        //检查是否符合avl树
        while (parent != null) {

            if (parent.bf == 2) {
                if (node.bf == -1) {
                    //右左双旋
                    turnRL(parent);
                } else {
                    //左旋
                    turnL(parent);

                }
                break;
            } else if (parent.bf == -2) {
                if (node.bf == -1) {
                    //右旋
                    turnR(parent);

                } else {
                    //左右双旋
                    turnLR(parent);
                }
                break;
            }

            if (parent.bf == 0) {
                break;
            }
            //修改当前节点父节点的平衡因子
            if (parent == root) {
                break;
            }
            node = parent;
            parent = node.parent;
            if (parent.left == node) {
                parent.bf -= 1;
            } else {
                parent.bf += 1;
            }
        }
    }

    //左旋
    public void turnL(TreeNode parent) {
        TreeNode node = parent.right;
        TreeNode nodeL = node.left;
        parent.right = nodeL;
        if (nodeL != null) {
            nodeL.parent = parent;
        }
        node.parent = parent.parent;
        parent.parent = node;
        node.left = parent;
        parent.bf = 0;
        node.bf = 0;
        if (parent == root) {
            root = node;
        } else {
            if (node.parent.left == parent) {
                node.parent.left = node;
            } else {
                node.parent.right = node;
            }
        }
    }
    //右旋
    public void turnR(TreeNode parent) {
        TreeNode node = parent.left;
        TreeNode nodeR = node.right;
        parent.left = nodeR;
        if (nodeR != null) {
            nodeR.parent = parent;
        }
        node.parent = parent.parent;
        parent.parent = node;
        node.right = parent;
        parent.bf = 0;
        node.bf = 0;
        if (parent == root) {
            root = node;
        } else {
            if (node.parent.left == parent) {
                node.parent.left = node;
            } else {
                node.parent.right = node;
            }
        }
    }
    //左右双旋
    public void turnLR(TreeNode parent) {
        TreeNode node = parent.left;
        TreeNode nodeR = node.right;
        int bf = nodeR.bf;
        turnL(parent.left);
        turnR(parent);
        if (bf == 1) {
            node.bf = -1;
            nodeR.bf = 0;
            parent.bf = 0;
        } else if (bf == -1) {
            node.bf = 0;
            nodeR.bf = 0;
            parent.bf = 1;
        }
    }
    //右左双旋
    public void turnRL(TreeNode parent) {
        TreeNode node = parent.right;
        TreeNode nodeL = node.left;
        int bf = nodeL.bf;
        turnR(parent.right);
        turnL(parent);
        if (bf == 1) {
            node.bf = 0;
            nodeL.bf = 0;
            parent.bf = -1;
        } else if (bf == -1) {
            node.bf = 1;
            nodeL.bf = 0;
            parent.bf = 0;
        }
    }
}
